import pandas as pd
from augmentation import augment_dataframe_to_target

# 读取原始数据
df = pd.read_csv('datasets/ddr/DR_grading.csv')

# 分离0级和1级数据
df_0 = df[df['dr_grade'] == 0].copy()
df_1 = df[df['dr_grade'] == 1].copy()

# 对1级数据进行增强到1000张（如果不足）
df_1 = augment_dataframe_to_target(
    df_1,
    label_col='dr_grade',
    target_count=1000,
    img_dir='datasets/ddr/images',
    save_dir='datasets/ddr/images'
)

# 分别采样1000张
df_0_sampled = df_0.sample(n=1000, random_state=42)
df_1_sampled = df_1.sample(n=1000, random_state=42)

# 合并并打乱
df_sampled = pd.concat([df_0_sampled, df_1_sampled], ignore_index=True)
df_sampled = df_sampled.sample(frac=1, random_state=42).reset_index(drop=True)

# 保存
df_sampled.to_csv('datasets/ddr/DR_grading_sampled.csv', index=False)

print(f"采样完成！共 {len(df_sampled)} 条数据")
print(f"0级: {len(df_0_sampled)} 条")
print(f"1级: {len(df_1_sampled)} 条 (增强后总数: {len(df_1)})")
